#include <assert.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <cusolverDn.h>

#include "pdm.h"
#include <Utility/cuda_device.cuh>
#include "solver_kernels_cuda.cuh"

#include "helper.hpp"
namespace shermit {
namespace fitting {
namespace curve {
    /*solve PDM least square Ap*x=b using QR factorization, where
	  Ap is nd x nc,
	  x is nc x 2,
	  b is nd x 2
	*/
	void CudaPDMSolverLU64::solve(double error_target) {
		/* step 1: initialization */
		A.resize(nd * nc);
		tau.resize(nc);
		pt_soa.resize(nd * 2);
		const int m = nd;
		const int n = nc - 3;
		const int lda = nd;
		const int ldb = nd;
		const int nrhs = 2;
		double* bfm_ptr = getRawPtr(A);  // base function matrix
		double* A_ptr = getRawPtr(A);
		double* b_ptr = getRawPtr(pt_soa);
		double* tau_ptr = getRawPtr(tau);

		const nv::vec2<double>* pts_ptr = projector.getDataPoints();
		auto& CudaDevice = MCudaDevice::getInstance();

		cublasHandle_t cublas_h = NULL;
		cusolverDnHandle_t cusolver_h = NULL;
		cusolverDnCreate(&cusolver_h);
		cublasCreate(&cublas_h);
		
		/* step 2: query and allocate workspace */
		int lwork_geqrf, lwork_ormqr;
		cusolverStatus_t status;
		status = cusolverDnDgeqrf_bufferSize(cusolver_h, m, n, A_ptr, lda, &lwork_geqrf);
		assert(status == CUSOLVER_STATUS_SUCCESS);
		status = cusolverDnDormqr_bufferSize(cusolver_h, CUBLAS_SIDE_LEFT, CUBLAS_OP_T,
			m, nrhs, n, A_ptr, lda, tau_ptr, b_ptr, ldb, &lwork_ormqr);
		assert(status == CUSOLVER_STATUS_SUCCESS);
		int lwork = (lwork_geqrf > lwork_ormqr) ? lwork_geqrf : lwork_ormqr;
		workspace.resize(lwork);
		thrust::device_vector<int> info(1, 0);
		double* buffer = getRawPtr(workspace);
		int* info_ptr = getRawPtr(info);

		int _iter = 0;
		double* t_ptr;
		double rmse;
		while (_iter < max_iter_solver) {
			projector.computeProjection(curve, max_iter_projection, spr, &rmse, &t_ptr);
			if (rmse < error_target) break;

			/* load data points */
			auto gridblk = CudaDevice.getOptLaunchParams(nd, 0, calcPtSOA64);
			calcPtSOA64 << <gridblk.first, gridblk.second >> >
				(nd, pts_ptr, getRawPtr(pt_soa));
			cudaDeviceSynchronize();
			getLastCudaError("calcPtSOA64 failure");

			gridblk = CudaDevice.getOptLaunchParams(lda*n, 0, calcBaseFuncMatrixClosed64);
			calcBaseFuncMatrixClosed64 << <gridblk.first, gridblk.second >> >
				(curve, lda, n, t_ptr, bfm_ptr);
			cudaDeviceSynchronize();
			getLastCudaError("calcBaseFuncMatrix64 failure");
			if (0) {
				thrust::host_vector<double> h_A(nd * nc);
				h_A.assign(A.begin(), A.end());
				printMatrix(h_A.data(), m, n, lda, "A");
			}

			/* step 3: compute QR factorization */
			status = cusolverDnDgeqrf(cusolver_h, m, n, A_ptr, lda, tau_ptr, buffer, lwork, info_ptr);
			cudaDeviceSynchronize();
			assert(status == CUSOLVER_STATUS_SUCCESS);
			/* check if QR is good or not */
			int h_info = 1;
			checkCudaErrors(GPU2CPU(&h_info, info_ptr, sizeof(int)));
			if(h_info != 0) printf("solver iter %d: after geqrf: info = %d\n", _iter, h_info);
			assert(h_info == 0);

			/* step 4: compute Q^T*B */
			status = cusolverDnDormqr(
				cusolver_h, CUBLAS_SIDE_LEFT, CUBLAS_OP_T,
				m, nrhs, n, A_ptr, lda, tau_ptr, b_ptr, ldb, buffer, lwork, info_ptr);
			checkCudaErrors(cudaDeviceSynchronize());
			assert(status == CUSOLVER_STATUS_SUCCESS);
			/* check if QR is good or not */
			checkCudaErrors(GPU2CPU(&h_info, info_ptr, sizeof(int)));
			if (h_info != 0) printf("solver iter %d: after ormqr: info = %d\n", _iter, h_info);
			assert(h_info == 0);
			if (0) {
				thrust::host_vector<double> h_b(pt_soa.size());
				h_b.assign(pt_soa.begin(), pt_soa.end());
				printMatrix(h_b.data(), m, nrhs, ldb, "Q^T*B");
			}


			/* step 5: compute x = R \ Q^T*B */
			const double alpha = 1;
			auto blas_status = cublasDtrsm(
				cublas_h, CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT,
				n, nrhs, &alpha, A_ptr, lda, b_ptr, ldb);
			cudaDeviceSynchronize();
			assert(blas_status == CUBLAS_STATUS_SUCCESS);
			if (0) {
				thrust::host_vector<double> h_b(pt_soa.size());
				h_b.assign(pt_soa.begin(), pt_soa.end());
				printMatrix(h_b.data(), n, 2, ldb, "x = R \ Q^T*B");
			}
			
			gridblk = CudaDevice.getOptLaunchParams(nc, 0, setPts64);
			setPts64 << <gridblk.first, gridblk.second >> >
				(nc, n, b_ptr, b_ptr + ldb, curve.control_pts);
				
			cudaDeviceSynchronize();
			getLastCudaError("setPts64 failure");
			_iter++;
		}
		printf("PDM Solver iter: %d, RMSE: %lf\n", _iter, rmse);
		cusolverDnDestroy(cusolver_h);
		cublasDestroy(cublas_h);
		return;
	}

}
}
}